#include <iostream>

#include <memory>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <sys/types.h>

#ifdef ENABLE_CNC
#include "./cnc/CNCTask.hpp"
#endif

#ifdef ENABLE_SWARM
#include "./swarm/SWARMTask.hpp"
#endif

#include "./base/base_uday_wrapper.hpp"

#include "./support/Threading.hpp"
#include "./support/check.hpp"

#include "./common/definitions.h"
#include "./common/Runnable.hpp"

#ifdef ENABLE_PAPI
#include "./support/PapiHook.hpp"
#endif

ostream &operator<<(ostream &os, Runnable &r) { 
  r.print(os);
  return os;
}

extern "C" {
extern void heat_3d(double M[][2][Npadded][Npadded][Npadded], int pT, int pN);
}

// The data
static int Titer = 0, Niter = 0;

#ifndef STATIC_ALLOC
double (* A)[2][Npadded][Npadded][Npadded];
double (* Acheck)[2][Npadded][Npadded][Npadded];
#else
double A[1][2][Npadded][Npadded][Npadded];
double Acheck[1][2][Npadded][Npadded][Npadded];
#endif // STATIC_ALLOC

//
//  Test whether a and b are equal (where epsilon is the error).
//  I.e.,
//
//  |a-b| <= epsilon |a|   or |a-b| <= epsilon |b|
//

// Checks the data 
int check() {
  int i;

  std::cout << "START CHECKING RESULTS" << std::endl;

  heat_3d(Acheck, Titer, Niter);

  for (i=0; i<Niter; i++) {
    if (check_matrices("A0", "A0check", 
		       &A[0][0][i][0][0], 
		       &Acheck[0][0][i][0][0], 
		       Niter, Npadded))
      return 1;
    if (check_matrices("A1", "A1check", 
		       &A[0][1][i][0][0], 
		       &Acheck[0][1][i][0][0], 
		       Niter, Npadded))
      return 1;
  }

  return 0;
}


// Initializing the data
void initialize_once() {
  int i, j, k;

  int SLABS = 2;
  int ARRAY_SIZE = Npadded * Npadded * Npadded;

#ifndef STATIC_ALLOC
  A = (double (*)[2][Npadded][Npadded][Npadded]) malloc (SLABS * ARRAY_SIZE * sizeof(double));
  Acheck = (double (*)[2][Npadded][Npadded][Npadded]) malloc (SLABS * ARRAY_SIZE * sizeof(double));
#endif // STATIC_ALLOC

  double d = 0.0; // used for template instantiation
  Threading::pinPages(A, SLABS, ARRAY_SIZE, d);

  // Rand not thread safe, separate from init of A
  for (i = 0; i < Npadded; i++) {
    for (j = 0; j < Npadded; j++) {
      for (k = 0; k < Npadded; k++) {
	Acheck[0][0][i][j][k] = 0.0;//(double)(rand() / ((double)RAND_MAX + 1));
	Acheck[0][1][i][j][k] = 0.0;//(double)(rand() / ((double)RAND_MAX + 1));
      }
    }
  }

  // Just set 10000 random values in each array
  for (i=0; i<10000; i++) {
    Acheck[0][0][rand() % Npadded][rand() % Npadded][rand() % Npadded] =
      1000.0 * (double)(rand() / ((double)RAND_MAX + 1));
    Acheck[0][1][rand() % Npadded][rand() % Npadded][rand() % Npadded] =
      1000.0 * (double)(rand() / ((double)RAND_MAX + 1));
  }

  for (i = 0; i < Npadded; i++) {
    for (j = 0; j < Npadded; j++) {
      for (k = 0; k < Npadded; k++) {
	A[0][0][i][j][k] = Acheck[0][0][i][j][k];
	A[0][1][i][j][k] = Acheck[0][1][i][j][k];
      }
    }
  }
}


int main (int argc, char *argv[])
{
  bool perform_check = false, papi = false;
  int argi;
  int num_trials = 1;
  int code_version = 0;

  // Command line: stencil [-v] [-c] [-T num_trial] [-V code_version] -i Titer Niter
  if (argc < 3 || argc > 10) {
    fprintf(stderr, "Incorrect number of arguments\n");
    fprintf(stderr, "Usage: stencil [-v] [-c] [-T num_trial] \
                         [-V code_version] -i Titer Niter\n");
    return -1;
  }

  for (argi=1; argi<argc; argi++) {
    if (strcmp("-c", argv[argi]) == 0) {
      perform_check = true;
      continue;
    }
    if (strcmp("-T", argv[argi]) == 0) {
      argi++;
      num_trials = atol(argv[argi]);
      std::cout << num_trials << " trials\n";
      continue;
    }
    if (strcmp("-P", argv[argi]) == 0) {
      papi = true;
      std::cout << "Use PAPI\n";
      continue;
    }
    if (strcmp("-V", argv[argi]) == 0) {
      argi++;
      code_version = atol(argv[argi]);
      continue;
    }

    break;
  }


  if (strcmp("-i", argv[argi]) != 0) {
    fprintf(stderr, "Expecting -i problem_size");
    fprintf(stderr, "Usage: stencil [-v] [-c] [-T num_trial] \
                       [-V code_version] -i Titer Niter\n");
    exit(1);
  }

  argi++;
  Titer = atol(argv[argi++]);
  Niter = atol(argv[argi++]);
    
  if (argi < argc) {      
    fprintf(stderr, "Unexpected parameters");
    fprintf(stderr, "Usage: stencil [-v] [-c] [-T num_trial] \
                       [-V code_version] -i Titer Niter\n");
  }

  // Initialize arrays before wrapper
  initialize_once();


  { // Begin run
    std::shared_ptr<Runnable> s_tmp;

    if (code_version == 0) {
#ifdef ENABLE_CNC
      CNCTask<0, 4, TASKTAGDIMMACRO, INTERTASKDIMMACRO> *st = 
        new CNCTask<0, 4, TASKTAGDIMMACRO, INTERTASKDIMMACRO>(A, Titer, Niter, Niter, Niter);
      s_tmp = std::shared_ptr<CNCTask<0, 4, TASKTAGDIMMACRO, INTERTASKDIMMACRO> >(st);
#else
      std::cerr << "\n!!! ERROR: V 1 statically linked to CNC but CNC compilation not enabled\nEXITING\n";
      exit(1);
#endif
    } 

    if (code_version == 1) {
#ifdef ENABLE_SWARM
      SWARMTask<0, 4, TASKTAGDIMMACRO, INTERTASKDIMMACRO> *st = 
        new SWARMTask<0, 4, TASKTAGDIMMACRO, INTERTASKDIMMACRO>(A, Titer, Niter, Niter, Niter);
      s_tmp = std::shared_ptr<SWARMTask<0, 4, TASKTAGDIMMACRO, INTERTASKDIMMACRO> >(st);
#else
      std::cerr << "\n!!! ERROR: V 1 statically linked to SWARM but SWARM compilation not enabled\nEXITING\n";
      exit(1);
#endif
    } 

    if (code_version > 1) {
      Uday *ud = new Uday(A, Titer, Niter, Niter, Niter, code_version);
      s_tmp = std::shared_ptr<Uday>(ud);
    }    
      
    assert(s_tmp.get() != nullptr);
    Runnable &s = *s_tmp;
    std::cout << "Running: " << s << std::endl;
    
#ifdef ENABLE_PAPI
    PAPI<Runnable> p;
    if (papi) {
      p.papi_loop(s);
    } else {
      s();
    }
#else
    for (int t=0; t<num_trials; t++) {
      s();
    }
#endif

    std::cout << "DONE\n";
  } // End run


  if (perform_check) {
    check();
  }

#ifndef STATIC_ALLOC
  free(A);
  free(Acheck);
#endif // STATIC_ALLOC

  return 0;
}
